import asyncio
from loguru import logger
from typing import Optional

from app.models.task.task_data import Task
from ..sse_manager import sse_manager


class TaskMonitor:
    """任务监控器"""

    def __init__(self):
        self._monitoring_tasks: dict = {}
        self._monitor_lock = asyncio.Lock()

    async def start_monitoring(self, task: Task, timeout: Optional[int] = None):
        """开始监控任务"""
        async with self._monitor_lock:
            if timeout:
                # 创建超时监控任务
                monitor_task = asyncio.create_task(self._timeout_monitor(task, timeout))
                self._monitoring_tasks[task.id] = monitor_task

    async def stop_monitoring(self, task_id: str):
        """停止监控任务"""
        async with self._monitor_lock:
            if task_id in self._monitoring_tasks:
                monitor_task = self._monitoring_tasks[task_id]
                if not monitor_task.done():
                    monitor_task.cancel()
                del self._monitoring_tasks[task_id]

    async def _timeout_monitor(self, task: Task, timeout: int):
        """超时监控"""
        try:
            await asyncio.sleep(timeout)
            logger.warning(f"任务超时: {task.name} (ID: {task.id})")

            # 发送超时通知
            await sse_manager.send_system_message(
                f"任务 {task.name} 执行超时 ({timeout}秒)", level="warning"
            )

        except asyncio.CancelledError:
            # 任务正常完成，取消监控
            pass

    async def cleanup(self):
        """清理所有监控任务"""
        async with self._monitor_lock:
            for monitor_task in self._monitoring_tasks.values():
                if not monitor_task.done():
                    monitor_task.cancel()
            self._monitoring_tasks.clear()


# 全局任务监控器
task_monitor = TaskMonitor()
